from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
import isaacgym
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_vecenv_stable_baselines import PCGVecEnvStableBaselines

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback


LOGDIR = "./training_logs"

trial_dir = "./trial2/"
# trial_dir = ""


total_generations = 22
finess_function = 2
test_trail = 14

# load performance.json and convert to a float list
def load_performance(pth):
    performance = None
    with open(pth, "r") as f:
        for line in f:
            performance = line
    performance = eval(performance)
    performance = np.array(performance)
    return performance

# compute performance
def compute_performance(performance_records, verbose=True):

    if verbose:
        print("---------------------")

    evaluations = []
    for i in range(len(performance_records)):

        larva_eval = performance_records[i][0]
        adult_eval = performance_records[i][1]

        # 3.1. maximum performance drop on old maps
        sr_drop_max = -1000
        for j in range(len(larva_eval)-1):
            sr_drop = larva_eval[j] - adult_eval[j]
            if sr_drop > sr_drop_max:
                sr_drop_max = sr_drop

        # 3.2. performance gain on new maps
        sr_gain = adult_eval[-1] - larva_eval[-1]

        # 3.3. success rate on new maps
        sr_new = adult_eval[-1]

        if verbose:
            print("-------------")
            print("sr_initial: ", larva_eval[-1], " i : ", i)
            print("sr_final: ", sr_new, " i : ", i)

        old_sr_threshold = 0.73
        new_sr_threshold = 0.80

        evaluations.append(copy.deepcopy([sr_drop_max, sr_gain, sr_new]))

    scores = []
    for i in range(len(evaluations)):

        adult_eval = performance_records[i][1]
        sr_new = adult_eval[-1]

        min_old_sr = 10000
        for j in range(len(adult_eval)-1):
            if adult_eval[j] < min_old_sr:
                min_old_sr = adult_eval[j]

        if min_old_sr < old_sr_threshold or sr_new < new_sr_threshold:
            if verbose:
                print("min_old_sr : ", min_old_sr, " sr_new : ", sr_new)
            continue

        if finess_function ==1:
            score = sqrt((evaluations[i][0]+1)**2 + (evaluations[i][1]-1)**2 + (evaluations[i][2]-1)**2)
        elif finess_function ==2:
            score = abs(evaluations[i][0]+1) + abs(evaluations[i][1]-1) + abs(evaluations[i][2]-1)
        scores.append((i, copy.deepcopy(score)))

    scores = sorted(scores, key=lambda x: x[1])

    return copy.deepcopy(scores), copy.deepcopy(evaluations)


def switch_element(evaluations,index1,index2,item):

    evaluations_ = copy.deepcopy(evaluations)

    temp = evaluations_[index1][item]
    evaluations_[index1][item] = evaluations_[index2][item]
    evaluations_[index2][item] = temp

    return evaluations_

# construct string:performance_records_gen_* where * ranges from 0 to 16
def construct_performance_records_gen_path(gen):
    return trial_dir+"performance_records_gen_" + str(gen) + ".json"

pth_list = []

for g in range(0, total_generations):
    for d in range(0, 8):
        pth_list.append(trial_dir+"gen_" + str(g) + "_dec_" + str(d) + ".json")
        print(pth_list[-1])

l_id = 0

# m_env = PCGVecEnvStableBaselines(seed_pth = pth_list[l_id], headless_ = False)
m_env = PCGVecEnvStableBaselines(headless_ = False)

observation = m_env.reset()

print("num_envs:",m_env.num_envs)

all_obs = np.zeros((16, 84, 84, 4), dtype=np.uint8)

all_actions = np.zeros(m_env.num_matrix_envs)
for i in range(m_env.num_matrix_envs):
    all_actions[i] = -2

steps_ = 0
reset_id = 0

PCGWorker_ = PCGWorker(9,9)

manual = False

m_env.pause()

# load performance records
performance_records_all = []
for g in range(total_generations):

    pr_pth = construct_performance_records_gen_path(g)
    performance_records = load_performance(pr_pth)
    performance_records_all.append(performance_records)

    for i in range(len(performance_records)):
        larva_eval = performance_records[i][0]
        adult_eval = performance_records[i][1]

        # print("------------- generation : ", g, " i : ", i)
        # print("sr_final: ", adult_eval[-1])

while True:

    # steps_ += 1

    observation, reward, done, info = m_env.step(all_actions)

    obs1 = observation[0]
    # transpose obs1 from (3,84,84) to (84,84,3)
    obs1 = np.transpose(obs1, (1,2,0))

    # display img
    cv2.imshow('img', obs1)
    key_ = cv2.waitKey(1)
    if key_ == 119:  # w
        action = 0
    elif key_ == 115:    # s
        action = 1
    elif key_ == 97:     # a
        action = 2
    elif key_ == 100:    # d
        action = 3
    elif key_ == 107:    # ccw
        action = 4
    elif key_ == 108:    # cw
        action = 5
    elif key_ == 105:    # reset
        m_env.reset()
    elif key_ == 114:    # resume
        m_env.resume()
        manual = True
    elif key_ == 112:    # pause
        m_env.pause()
    elif key_ == 83:    # change landscape
        l_id += 1
        if l_id >= len(pth_list):
            l_id = 0
        seed = PCGWorker_.from_file(pth_list[l_id])
        m_env.set_landscape(env_id = 0,seed_ = seed)
        # transform the number between 'gen_' and '_dec_' in the file name to int
        g = int(pth_list[l_id][pth_list[l_id].find("gen_")+4:pth_list[l_id].find("_dec_")])
        # transform the number after '_dec_' in the file name to int
        d = int(pth_list[l_id][pth_list[l_id].find("_dec_")+5:pth_list[l_id].find(".json")])
        adult_sr = performance_records_all[g][d][1][-1]
        print("landscape changed to:", pth_list[l_id], " adult_sr:", adult_sr)
    elif key_ == 81:
        l_id -= 1
        if l_id < 0:
            l_id = len(pth_list) - 1
        seed = PCGWorker_.from_file(pth_list[l_id])
        m_env.set_landscape(env_id = 0,seed_ = seed)
        # transform the number between 'gen_' and '_dec_' in the file name to int
        g = int(pth_list[l_id][pth_list[l_id].find("gen_")+4:pth_list[l_id].find("_dec_")])
        # transform the number after '_dec_' in the file name to int
        d = int(pth_list[l_id][pth_list[l_id].find("_dec_")+5:pth_list[l_id].find(".json")])
        adult_sr = performance_records_all[g][d][1][-1]
        print("landscape changed to:", pth_list[l_id], " adult_sr:", adult_sr)
    elif key_ == 109:    # m
        seed = PCGWorker_.generate()
        m_env.set_landscape(env_id = 0,seed_ = seed)
    else:
        action = -2

    # print(key_)

    for i in range(m_env.num_matrix_envs):
        all_actions[i] = action

    m_env.render()
